#-*- coding: utf-8 -*-
import math
import numpy as np
import matplotlib.pyplot as plt

import pat.dynamic_model as dm
import pat.frames as fr
import pat.utils as pu

sv_ax    = 0
sv_ay    = 1
sv_az    = 2
sv_p     = 3
sv_q     = 4
sv_r     = 5
sv_x     = 6
sv_y     = 7
sv_z     = 8
sv_phi   = 9
sv_theta = 10
sv_psi   = 11
sv_vx    = 12
sv_vy    = 13
sv_vz    = 14
sv_mx    = 15
sv_my    = 16
sv_mz    = 17
sv_size  = 18

iv_pitch = 0
iv_roll  = 1
iv_col   = 2
iv_yaw   = 3
iv_size  = 4
    
class Param():
    def __init__(self):
        self.wind_gust_max = np.array([0., 0., 0.])

    def __str__(self):
        return "wind gust max:{:s}".format(self.wind_gust_max)


class DynamicModel(dm.DynamicModel):

    sv_ax    = 0
    sv_ay    = 1
    sv_az    = 2
    sv_p     = 3
    sv_q     = 4
    sv_r     = 5
    sv_x     = 6
    sv_y     = 7
    sv_z     = 8
    sv_phi   = 9
    sv_theta = 10
    sv_psi   = 11
    sv_vx    = 12
    sv_vy    = 13
    sv_vz    = 14
    sv_mx    = 15
    sv_my    = 16
    sv_mz    = 17
    sv_size  = 18

    iv_pitch = 0
    iv_roll  = 1
    iv_col   = 2
    iv_yaw   = 3
    iv_size  = 4
    
    def __init__(self, args=None):
        dm.DynamicModel.__init__(self)
        self.X = np.zeros(DynamicModel.sv_size)
        self.c_dm = __import__("libpat_dm_hk")
        self.p = Param()

    def name(self):
        return "rotorcraft Hudson/Kahn" 

    def reset(self, X0=None):
        if X0<>None: self.X = X0
        else: self.X = np.zeros(DynamicModel.sv_size)

    def trim(self, args=None):
        Xe = np.zeros(DynamicModel.sv_size); Xe[sv_z] = -4.
        Ue = np.array([0.0,  0.0,  0.16024332,  0.10031169])
        return Xe, Ue

    def run(self, dt, U):
        self.c_dm.run(self.X, dt, U)
        return self.X

    def param(self):
        return str(self.p)

    def state_str(self):
        return """pos: {:-.2f}, {:-.2f}, {:-.2f} m
        vel: {:-.2f}, {:-.2f}, {:-.2f} m/s
        att:    {:-.2f}, {:-.2f}, {:-.2f} deg
        """.format(self.X[sv_x], self.X[sv_y], self.X[sv_z],
                   self.X[sv_vx], self.X[sv_vy], self.X[sv_vz],
                   pu.deg_of_rad(self.X[sv_phi]), pu.deg_of_rad(self.X[sv_theta]), pu.deg_of_rad(self.X[sv_psi]))


    def state_SixDOFfEuclidianEuler(self):
        X = np.zeros(fr.SixDOFfEuclidianEuler.size)
        X[fr.SixDOFfEuclidianEuler.x:fr.SixDOFfEuclidianEuler.z+1] = self.X[sv_x:sv_z+1]
        X[fr.SixDOFfEuclidianEuler.phi:fr.SixDOFfEuclidianEuler.psi+1] = self.X[sv_phi:sv_psi+1]
        return X

    def plot_trajectory(self, time, X, U=None):
        ax = plt.subplot(3,3,1)         # x
        plt.plot(time, X[:,self.sv_x])
        pu.decorate(ax, title="$x$")
        ax=plt.subplot(3,3,2)           # y
        plt.plot(time, X[:,self.sv_y])
        pu.decorate(ax, title="$y$")
        ax=plt.subplot(3,3,3)           # z
        plt.plot(time, X[:,8]) 
        pu.decorate(ax, title="$z$")

        ax=plt.subplot(3,3,4)
        plt.plot(time, pu.deg_of_rad(X[:,self.sv_phi]))  # phi
        pu.decorate(ax, title="$\phi$")
        ax=plt.subplot(3,3,5)
        plt.plot(time, pu.deg_of_rad(X[:,self.sv_theta])) # theta
        pu.decorate(ax, title="$\\theta$")
        ax=plt.subplot(3,3,6)
        plt.plot(time, pu.deg_of_rad(X[:,self.sv_psi])) # psi
        pu.decorate(ax, title="$\psi$")
        
        if U<>None:
            ax = plt.subplot(3,3,7)
            plt.plot(time, U[:,0]) # pitch
            plt.plot(time, U[:,1]) # pitch
            pu.decorate(ax, title="$pitch/roll$", legend=["pitch", "roll"])
            ax = plt.subplot(3,3,8)
            plt.plot(time, U[:,2]) # col
            pu.decorate(ax, title="$col$")
            ax = plt.subplot(3,3,9)
            plt.plot(time, U[:,3]) # yaw
            pu.decorate(ax, title="$yaw$")
            
